from Helper import *
from Model import *

import numpy as np
            
plateName=r"#京沪津渝冀晋蒙辽吉黑苏浙皖闽赣鲁豫鄂湘粤桂琼川贵云藏陕甘青宁新学警港澳挂使领民航危0123456789ABCDEFGHJKLMNPQRSTUVWXYZ险品"
detect_model = load_yolov8_model(r'best.pt', torch.device("cuda" if torch.cuda.is_available() else "cpu"))
plate_rec_model = init_ocr_model(torch.device("cuda" if torch.cuda.is_available() else "cpu"), r'plate_rec_color.pth', plateName, is_color=True)

def process_img(img: np.ndarray, detect_model  = detect_model, plate_rec_model = plate_rec_model, img_size=640, conf_threshlod=0.3, iou_threshold=0.5, device=torch.device("cuda" if torch.cuda.is_available() else "cpu")) -> np.ndarray:
    return det_rec_plate(img, detect_model,plate_rec_model, img_size, conf_threshlod, iou_threshold, device)